import os
import numpy as np
import torch
from torch import nn, optim
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt

from model import build_resnet_18_model  # 可根据需要更换模型

# 全局变量
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # 禁用 GPU
DEVICE = torch.device('cpu')
BATCH_SIZE = 32
EPOCHS = 100
PATIENCE = 10  # 早停 patience
MODEL_SAVE_PATH = '../weights/cough_recognition_model.pth'
BEST_MODEL_SAVE_PATH = '../weights/best_model.pth'
HISTORY_SAVE_PATH = '../weights/history.npy'


def load_data():
    train_data = np.load('../data/train_data.npy')  # 原始形状 (N, H, W, 1)
    val_data = np.load('../data/val_data.npy')
    train_labels = np.load('../data/train_labels.npy')
    val_labels = np.load('../data/val_labels.npy')
    # 转换为 (N, 1, H, W)
    train_data = np.transpose(train_data, (0, 3, 1, 2))
    val_data = np.transpose(val_data, (0, 3, 1, 2))
    return train_data, val_data, train_labels, val_labels


def create_dataloaders(train_data, val_data, train_labels, val_labels):
    train_dataset = TensorDataset(torch.tensor(train_data, dtype=torch.float32),
                                  torch.tensor(train_labels, dtype=torch.long))
    val_dataset = TensorDataset(torch.tensor(val_data, dtype=torch.float32),
                                torch.tensor(val_labels, dtype=torch.long))
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
    return train_loader, val_loader


def train_model(model, train_loader, val_loader, num_epochs, patience, device):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())

    best_val_loss = float('inf')
    epochs_no_improve = 0
    history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        running_corrects = 0
        total = 0
        for inputs, labels in train_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            running_corrects += (preds == labels).sum().item()
            total += inputs.size(0)
        epoch_loss = running_loss / total
        epoch_acc = running_corrects / total

        # 验证阶段
        model.eval()
        val_loss = 0.0
        val_corrects = 0
        val_total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * inputs.size(0)
                _, preds = torch.max(outputs, 1)
                val_corrects += (preds == labels).sum().item()
                val_total += inputs.size(0)
        val_epoch_loss = val_loss / val_total
        val_epoch_acc = val_corrects / val_total

        history['train_loss'].append(epoch_loss)
        history['train_acc'].append(epoch_acc)
        history['val_loss'].append(val_epoch_loss)
        history['val_acc'].append(val_epoch_acc)

        print(f"Epoch {epoch + 1}/{num_epochs}: train_loss={epoch_loss:.4f}, train_acc={epoch_acc:.4f}, "
              f"val_loss={val_epoch_loss:.4f}, val_acc={val_epoch_acc:.4f}")

        # 早停和最佳模型保存
        if val_epoch_loss < best_val_loss:
            best_val_loss = val_epoch_loss
            epochs_no_improve = 0
            torch.save(model.state_dict(), BEST_MODEL_SAVE_PATH)
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= patience:
                print("Early stopping")
                break

    return model, history


def save_model(model, history):
    torch.save(model.state_dict(), MODEL_SAVE_PATH)
    np.save(HISTORY_SAVE_PATH, history)


def plot_metrics(history):
    plt.figure(figsize=(12, 4))
    # Loss 曲线
    plt.subplot(1, 2, 1)
    plt.plot(history['train_loss'], label='Train Loss')
    plt.plot(history['val_loss'], label='Validation Loss')
    plt.title('Loss over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    # Accuracy 曲线
    plt.subplot(1, 2, 2)
    plt.plot(history['train_acc'], label='Train Accuracy')
    plt.plot(history['val_acc'], label='Validation Accuracy')
    plt.title('Accuracy over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.tight_layout()
    plt.show()


def train():
    train_data, val_data, train_labels, val_labels = load_data()
    train_loader, val_loader = create_dataloaders(train_data, val_data, train_labels, val_labels)
    # 获取输入形状，注意：train_data.shape = (N, 1, H, W)
    input_shape = train_data.shape[1:]  # (1, H, W)
    model = build_resnet_18_model(input_shape)
    model.to(DEVICE)
    model, history = train_model(model, train_loader, val_loader, EPOCHS, PATIENCE, DEVICE)
    save_model(model, history)
    plot_metrics(history)


if __name__ == '__main__':
    train()
